package test.demo;

import cn.hutool.core.io.resource.ResourceUtil;
import tech.waterism.modelbase.PredictBase;
import tech.waterism.modelbase.PredictResult;
import com.alibaba.fastjson.JSON;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.Arrays;
import java.util.Calendar;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 三水源蓄满产流模型
 * 系统模型集成
 *
 * @author 初京刚
 * @version 6.0
 */
public class SMS_3_TEST extends PredictBase {

    private double B;
    private double C;
    private double WUMx;
    private double WUM;
    private double WLMx;
    private double WLM;
    private double WDM;
    private double WM;
    private double WMM;
    private double KG;
    private double KI;
    private double SMM;
    private double SM;
    private double EX;
    private double IM;
    private double WUP;
    private double WLP;
    private double WDP;
    private double SP;
    private double FRP;
    private double K;
    private double[] ES;
    private double PE;
    private Date[] TM;
    private double[] PP;
    private double T;

    public SMS_3_TEST(Map<String, String> data) {
        super();

        IM = Double.parseDouble(data.get("IM"));
        B = Double.parseDouble(data.get("B"));
        K = Double.parseDouble(data.get("K"));
        KG = Double.parseDouble(data.get("KG"));
        KI = Double.parseDouble(data.get("KI"));
        C = Double.parseDouble(data.get("C"));
        WM = Double.parseDouble(data.get("WM"));
        SM = Double.parseDouble(data.get("SM"));
        WUMx = Double.parseDouble(data.get("WUMx"));
        WLMx = Double.parseDouble(data.get("WLMx"));
        EX = Double.parseDouble(data.get("EX"));
        ES = JSON.parseObject(data.get("ES"), double[].class);
        WUP = Double.parseDouble(data.get("WUP"));
        WLP = Double.parseDouble(data.get("WLP"));
        WDP = Double.parseDouble(data.get("WDP"));
        SP = Double.parseDouble(data.get("SP"));
        FRP = Double.parseDouble(data.get("FRP"));
        TM = JSON.parseObject(data.get("dt"), Date[].class);             // input[i].data.dt
        PP = JSON.parseObject(data.get("rain"), double[].class);  // input[i].data.rain
        T = Double.parseDouble(data.get("clen"));
    }

    public Map<String, double[]> SMS3model() {

        WUM = WUMx * WM;
        WLM = (1 - WUMx) * WLMx * WM;
        WDM = WM - WUM - WLM;
        double WP = WUP + WLP + WDP;

        double DIV = 5.0;
        WMM = (1.0 + B) * WM / (1.0 - IM);
        SMM = (1.0 + EX) * SM;

        int Pnum = PP.length;

        double[] WU = new double[Pnum + 1];
        double[] WL = new double[Pnum + 1];
        double[] WD = new double[Pnum + 1];
        double[] W = new double[Pnum + 1];
        double[] RSP = new double[Pnum];
        double[] RIP = new double[Pnum];
        double[] RGP = new double[Pnum];
        double[] RP = new double[Pnum];
        double[] S = new double[Pnum + 1];
        double[] FR = new double[Pnum + 1];
        WU[0] = WUP;
        WL[0] = WLP;
        WD[0] = WDP;
        W[0] = WP;
        S[0] = SP;
        FR[0] = FRP;

        if (KG + KI > 0.9) {
            double tmp = (KG + KI - 0.9) / (KG + KI);
            KG = KG - KG * tmp;
            KI = KI - KI * tmp;
        }

        double HGI = (1.0 - Math.pow((1.0 - KG - KI), (T / 24.0))) / (KG + KI);
        KG = HGI * KG;
        KI = HGI * KI;

        for (int i = 0; i < Pnum - 1; i++) {
            int IDAY = 31;
            Calendar time = Calendar.getInstance();
            time.setTime(TM[i]);
            Integer Mth = time.get(Calendar.MONTH) + 1; // 获得降雨开始月份
            switch (Mth) {
                case 4:
                case 6:
                case 9:
                case 11:
                    IDAY = 30;
                    break;
                case 2:
                    IDAY = 28;
                    break;
                default:
                    break;
            }
            double EM = ES[Mth - 1] / (IDAY * 24.0 / T);
            double EK = K * EM;
            PE = PP[i] - EK;
            if (PE >= 2 * DIV) {
                int ND = (int) Math.floor(PE / DIV);
                double[] PED = new double[ND];
                double[] RD = new double[ND];
                for (int II = 0; II < ND - 1; II++) {
                    PED[II] = DIV;
                }
                PED[ND - 1] = PE - (ND - 1) * DIV;

                Map<String, Double> WMap = YIELD1(EK, W[i], WU[i], WL[i], WD[i], RD, PED, ND);
                W[i + 1] = WMap.get("W");
                WU[i + 1] = WMap.get("WU");
                WL[i + 1] = WMap.get("WL");
                WD[i + 1] = WMap.get("WD");
                Map<String, Double> RMap = DIVI31(FR[i], S[i], RD, PED, ND);
                S[i + 1] = RMap.get("S");
                FR[i + 1] = RMap.get("FR");
                RSP[i] = RMap.get("RS");
                RIP[i] = RMap.get("RI");
                RGP[i] = RMap.get("RG");
            } else {
                int ND = 1;
                double[] PED = new double[ND];
                double[] RD = new double[ND];
                PED[0] = PE;
                Map<String, Double> map = YIELD1(EK, W[i], WU[i], WL[i], WD[i], RD, PED, ND);
                W[i + 1] = map.get("W");
                WU[i + 1] = map.get("WU");
                WL[i + 1] = map.get("WL");
                WD[i + 1] = map.get("WD");
                Map<String, Double> RMap = DIVI31(FR[i], S[i], RD, PED, ND);
                S[i + 1] = RMap.get("S");
                FR[i + 1] = RMap.get("FR");
                RSP[i] = RMap.get("RS");
                RIP[i] = RMap.get("RI");
                RGP[i] = RMap.get("RG");
            }
            if (RSP[i] < 0.0) {
                RSP[i] = 0.0;
            }
            if (RIP[i] < 0.0) {
                RIP[i] = 0.0;
            }
            if (RGP[i] < 0.0) {
                RGP[i] = 0.0;
            }
            RP[i] = RSP[i] + RIP[i] + RGP[i];
        }

        Map<String, double[]> simulationResult = new HashMap<>();

        simulationResult.put("RunoffSim", RP);
        simulationResult.put("RS", RSP);
        simulationResult.put("RI", RIP);
        simulationResult.put("RG", RGP);
        simulationResult.put("WU", WU);
        simulationResult.put("WL", WL);
        simulationResult.put("WD", WD);
        simulationResult.put("S", S);
        simulationResult.put("FR", FR);

        return simulationResult;
    }

    private Map<String, Double> YIELD1(double EK, double W, double WU, double WL, double WD, double[] RD, double[] PED,
                                       int ND) {
        if (PE <= 0.0) {
            double R = 0.0;
            if (WU + PE >= 0.0) {
                double EU = EK;
                double ED = 0.0;
                double EL = 0.0;
                WU = WU + PE;
            } else {
                double EU = WU + EK + PE;
                WU = 0.0;
                double EL = (EK - EU) * WL / WLM;
                if (WL < C * WLM) {
                    EL = C * (EK - EU);
                }
                if (WL - EL < 0.0) {
                    double ED = EL - WL;
                    EL = WL;
                    WL = 0.0;
                    WD = WD - ED;
                } else {
                    double ED = 0.0;
                    WL = WL - EL;
                }
            }
            W = WU + WL + WD;
        } else {
            double A = 0.0;
            if (WM - W < 0.0001) {
                A = WMM;
            } else {
                A = WMM * (1.0 - Math.pow((1.0 - W / WM), (1.0 / (1.0 + B))));
            }

            double R = 0.0;
            double PEDS = 0.0;
            for (int I = 0; I < ND; I++) {
                A = A + PED[I];
                PEDS = PEDS + PED[I];
                double RI = R;
                R = PEDS - WM + W;
                if (A < WMM) {
                    R = R + WM * Math.pow((1.0 - A / WMM), (1.0 + B));
                }
                RD[I] = R - RI;
            }
            double EU = EK;
            double EL = 0.0;
            double ED = 0.0;
            if (WU + PE - R <= WUM) {
                WU = WU + PE - R;
            } else {
                if (WU + WL + PE - R - WUM >= WLM) {
                    WU = WUM;
                    WL = WLM;
                    WD = W + PEDS - R - WU - WL;
                    if (WD > WDM) {
                        WD = WDM;
                    }
                } else {
                    WL = WU + WL + PE - R - WUM;
                    WU = WUM;
                }
            }
            W = WU + WL + WD;
        }
        Map<String, Double> WResult = new HashMap<>();
        WResult.put("W", W);
        WResult.put("WU", WU);
        WResult.put("WL", WL);
        WResult.put("WD", WD);
        return WResult;
    }

    private Map<String, Double> DIVI31(double FR, double S, double[] RD, double[] PED, int ND) {
        if (PE <= 0.0) {
            double RS = 0.0;
            double RG = S * KG * FR;
            double RI = S * KI * FR;
            S = S * (1.0 - KG - KI);

            Map<String, Double> RResult = new HashMap<>();
            RResult.put("FR", FR);
            RResult.put("S", S);
            RResult.put("RS", RS);
            RResult.put("RI", RI);
            RResult.put("RG", RG);
            return RResult;
        } else {
            double RB = IM * PE;
            double KID = (1.0 - Math.pow((1.0 - (KG + KI)), (1.0 / ND))) / (KG + KI);
            double KGD = KID * KG;
            KID = KID * KI;

            double RS = 0.0;
            double RI = 0.0;
            double RG = 0.0;

            for (int I = 0; I < ND; I++) {
                double TD = RD[I] - IM * PED[I];
                double X = FR;
                FR = TD / PED[I];
                S = X * S / FR;
                double RR = 0;
                if (S >= SM) {
                    RR = (PED[I] + S - SM) * FR;
                } else {
                    double AU = SMM * (1.0 - Math.pow((1.0 - S / SM), (1.0 / (1.0 + EX))));
                    if (AU + PED[I] < SMM) {
                        RR = (PED[I] - SM + S + SM * Math.pow((1.0 - (PED[I] + AU) / SMM), (1.0 + EX))) * FR;
                    } else {
                        RR = (PED[I] + S - SM) * FR;
                    }
                }
                RS = RR + RS;
                S = PED[I] - RR / FR + S;
                RG = S * KGD * FR + RG;
                RI = S * KID * FR + RI;
                S = S * (1.0 - KID - KGD);
            }
            RS = RS + RB;
            Map<String, Double> RResult = new HashMap<>();
            RResult.put("FR", FR);
            RResult.put("S", S);
            RResult.put("RS", RS);
            RResult.put("RI", RI);
            RResult.put("RG", RG);
            return RResult;
        }
    }

    private List<String> toStringList(double[] array0) {
        String[] array = new String[array0.length];
        for (int i = 0; i < array.length; i++) {
            array[i] = String.valueOf(array0[i]);
        }
        return Arrays.asList(array);
    }

    @Override
    public void init() {

    }

    @Override
    public PredictResult predict() {
        PredictResult result = new PredictResult();
        Map<String, double[]> map = SMS3model();
        double[][] RSim = new double[3][];
        //流域地表产流
        RSim[0] = map.get("RS");
        //流域土壤产流
        RSim[1] = map.get("RI");
        //流域地下产流
        RSim[2] = map.get("RG");

        //流域面积总产流 R
        result.setRunoffSim(map.get("RunoffSim"));

        result.setRSim(RSim);

        // 更新状态
        List<String> WUP = toStringList(map.get("WU"));
        List<String> WLP = toStringList(map.get("WL"));
        List<String> WDP = toStringList(map.get("WD"));
        List<String> SP = toStringList(map.get("S"));
        List<String> FRP = toStringList(map.get("FR"));
        Map<String, List<String>> status = new HashMap<>();
        //上层张力水含量
        status.put("WUP", WUP);
        //下层张力水含量
        status.put("WLP", WLP);
        //深层张力水含量
        status.put("WDP", WDP);
        //产流面积自由水深
        status.put("SP", SP);
        //产流面积系数
        status.put("FRP", FRP);
        result.setNewStatus(status);

        return result;
    }

    public static void main(String[] args) {
        // 从resources目录读取文件
        String a = ResourceUtil.readUtf8Str("20810200-方案178-SMS_3-1762739718330-20250605080000.txt");
        Map<String, String> data = JSON.parseObject(a, Map.class);
        SMS_3 model = new SMS_3(data);
        PredictResult result = model.predict();
        
        String outputPath = "output/SMS_3_output_chunjinggang.json";
        
        try {
            File outputFile = new File(outputPath);
            outputFile.getParentFile().mkdirs();
            try (FileWriter writer = new FileWriter(outputFile)) {
                writer.write(JSON.toJSONString(result));
            }
            System.out.println("结果已输出到: " + outputPath);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
